package gox

import (
	"sync"

	"gitee.com/burningsong/golib/pkg/errx"
	"github.com/zeromicro/go-zero/core/logx"
)

type Job func()

type Pool struct {
	jobs chan Job
	quit chan bool
	wait *sync.WaitGroup
}

func NewPool(maxGoroutines int) *Pool {
	pool := &Pool{
		jobs: make(chan Job, maxGoroutines),
		quit: make(chan bool),
		wait: &sync.WaitGroup{},
	}

	for i := 0; i < maxGoroutines; i++ {
		pool.wait.Add(1)
		go pool.worker()
	}

	return pool
}

func (p *Pool) worker() {
	defer func() {
		p.wait.Done()
		if r := recover(); r != nil {
			logx.Errorf("worker pool panic:%+v msg:%s", r, errx.GetStack())
		}
	}()
	for {
		select {
		case job := <-p.jobs:
			job()
		case <-p.quit:
			return
		}
	}
}

func (p *Pool) AddJob(job Job) {
	p.jobs <- job
}

func (p *Pool) Shutdown() {
	close(p.quit)
	p.wait.Wait()
}
